import sys
import hashlib
import os
import io
import threading
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                            QPushButton, QTreeWidget, QTreeWidgetItem, QFileDialog, QLabel, QMessageBox,
                            QProgressBar, QCheckBox, QGroupBox, QFormLayout, QSpinBox, QComboBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtGui import QPixmap, QImage
from PIL import Image, ImageQt
import rawpy
import send2trash

class ScanThread(QThread):
    progress_update = pyqtSignal(int, int)
    scan_complete = pyqtSignal(dict)
    
    def __init__(self, directory, min_size=0, max_size=float('inf')):
        super().__init__()
        self.directory = directory
        self.min_size = min_size
        self.max_size = max_size
        self.running = True
        
    def run(self):
        file_dict = {}
        total_files = 0
        processed_files = 0
        
        # 先计算文件总数
        for root, dirs, files in os.walk(self.directory):
            total_files += len(files)
        
        # 扫描文件
        for root, dirs, files in os.walk(self.directory):
            for file in files:
                if not self.running:
                    return
                    
                file_path = os.path.join(root, file)
                try:
                    file_size = os.path.getsize(file_path)
                    
                    # 应用文件大小过滤
                    if self.min_size <= file_size <= self.max_size:
                        file_hash = self.calculate_hash(file_path)
                        if file_hash in file_dict:
                            file_dict[file_hash].append(file_path)
                        else:
                            file_dict[file_hash] = [file_path]
                except (OSError, PermissionError):
                    pass
                    
                processed_files += 1
                self.progress_update.emit(processed_files, total_files)
        
        # 过滤只保留有重复的文件
        duplicate_dict = {k: v for k, v in file_dict.items() if len(v) > 1}
        self.scan_complete.emit(duplicate_dict)
    
    def calculate_hash(self, file_path, chunk_size=8192):
        hash_md5 = hashlib.md5()
        try:
            with open(file_path, 'rb') as f:
                while chunk := f.read(chunk_size):
                    if not self.running:
                        return None
                    hash_md5.update(chunk)
            return hash_md5.hexdigest()
        except (OSError, PermissionError):
            return None
            
    def stop(self):
        self.running = False

class FileDeduplicator(QMainWindow):
    def __init__(self):
        super().__init__()
        self.scan_thread = None
        self.initUI()
        
    def initUI(self):
        # 主窗口布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QVBoxLayout(main_widget)
        
        # 顶部控制面板
        control_panel = QHBoxLayout()
        
        # 目录选择
        self.btn_select = QPushButton("选择目录")
        control_panel.addWidget(self.btn_select)
        
        # 文件大小过滤
        size_filter_group = QGroupBox("文件大小过滤")
        size_filter_layout = QFormLayout()
        
        # 最小文件大小
        min_size_layout = QHBoxLayout()
        self.min_size_spin = QSpinBox()
        self.min_size_spin.setRange(0, 1000000)
        self.min_size_spin.setValue(0)
        self.min_size_unit = QComboBox()
        self.min_size_unit.addItems(["B", "KB", "MB", "GB"])
        self.min_size_unit.setCurrentIndex(1)  # 默认KB
        min_size_layout.addWidget(self.min_size_spin)
        min_size_layout.addWidget(self.min_size_unit)
        size_filter_layout.addRow("最小大小:", min_size_layout)
        
        # 最大文件大小
        max_size_layout = QHBoxLayout()
        self.max_size_spin = QSpinBox()
        self.max_size_spin.setRange(0, 1000000)
        self.max_size_spin.setValue(0)
        self.max_size_spin.setSpecialValueText("无限制")
        self.max_size_unit = QComboBox()
        self.max_size_unit.addItems(["B", "KB", "MB", "GB"])
        self.max_size_unit.setCurrentIndex(2)  # 默认MB
        max_size_layout.addWidget(self.max_size_spin)
        max_size_layout.addWidget(self.max_size_unit)
        size_filter_layout.addRow("最大大小:", max_size_layout)
        
        size_filter_group.setLayout(size_filter_layout)
        control_panel.addWidget(size_filter_group)
        
        # 添加到主布局
        main_layout.addLayout(control_panel)
        
        # 进度条
        progress_layout = QHBoxLayout()
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.status_label = QLabel("就绪")
        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.status_label)
        main_layout.addLayout(progress_layout)
        
        # 内容区域布局
        content_layout = QHBoxLayout()
        
        # 左侧文件列表
        left_panel = QVBoxLayout()
        self.tree_files = QTreeWidget()
        self.tree_files.setHeaderLabels(['文件名称', '文件路径', '文件大小'])
        self.tree_files.setColumnWidth(0, 200)
        self.tree_files.setColumnWidth(1, 300)
        left_panel.addWidget(self.tree_files)
        
        # 批量操作按钮
        batch_buttons = QHBoxLayout()
        self.btn_select_all = QPushButton("全选")
        self.btn_select_duplicates = QPushButton("智能选择")
        self.btn_delete = QPushButton("删除选中文件")
        batch_buttons.addWidget(self.btn_select_all)
        batch_buttons.addWidget(self.btn_select_duplicates)
        batch_buttons.addWidget(self.btn_delete)
        left_panel.addLayout(batch_buttons)
        
        # 右侧预览面板
        right_panel = QVBoxLayout()
        preview_group = QGroupBox("文件预览")
        preview_layout = QVBoxLayout()
        self.preview_label = QLabel()
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setMinimumSize(400, 400)
        self.preview_label.setStyleSheet("background-color: #f0f0f0; border: 1px solid #ddd;")
        preview_layout.addWidget(self.preview_label)
        preview_group.setLayout(preview_layout)
        right_panel.addWidget(preview_group)
        
        content_layout.addLayout(left_panel, 3)
        content_layout.addLayout(right_panel, 2)
        main_layout.addLayout(content_layout)
        
        # 连接信号槽
        self.btn_select.clicked.connect(self.select_directory)
        self.tree_files.itemClicked.connect(self.show_preview)
        self.btn_delete.clicked.connect(self.delete_selected)
        self.btn_select_all.clicked.connect(self.select_all_files)
        self.btn_select_duplicates.clicked.connect(self.select_duplicates)
        
        # 初始状态
        self.btn_delete.setEnabled(False)
        self.btn_select_all.setEnabled(False)
        self.btn_select_duplicates.setEnabled(False)
        
        # 窗口设置
        self.setGeometry(300, 300, 1000, 700)
        self.setWindowTitle('文件查重工具')
        
    def select_directory(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择查重目录")
        if dir_path:
            # 如果已有扫描线程在运行，先停止它
            if self.scan_thread and self.scan_thread.isRunning():
                self.scan_thread.stop()
                self.scan_thread.wait()
            
            # 获取文件大小过滤设置
            min_size = self.get_size_in_bytes(self.min_size_spin.value(), self.min_size_unit.currentText())
            max_size = float('inf')
            if self.max_size_spin.value() > 0:
                max_size = self.get_size_in_bytes(self.max_size_spin.value(), self.max_size_unit.currentText())
            
            # 清空当前结果
            self.tree_files.clear()
            self.preview_label.clear()
            
            # 更新UI状态
            self.btn_select.setText("停止扫描")
            self.progress_bar.setValue(0)
            self.status_label.setText("正在扫描...")
            
            # 创建并启动扫描线程
            self.scan_thread = ScanThread(dir_path, min_size, max_size)
            self.scan_thread.progress_update.connect(self.update_progress)
            self.scan_thread.scan_complete.connect(self.display_results)
            self.scan_thread.finished.connect(self.scan_finished)
            self.scan_thread.start()
        elif self.scan_thread and self.scan_thread.isRunning():
            # 如果点击时正在扫描，则停止扫描
            self.scan_thread.stop()
            self.scan_thread.wait()
            self.scan_finished()

    def get_size_in_bytes(self, size, unit):
        """将大小转换为字节"""
        if unit == "KB":
            return size * 1024
        elif unit == "MB":
            return size * 1024 * 1024
        elif unit == "GB":
            return size * 1024 * 1024 * 1024
        return size

    def update_progress(self, processed, total):
        """更新进度条"""
        percentage = int(processed / total * 100) if total > 0 else 0
        self.progress_bar.setValue(percentage)
        self.status_label.setText(f"已扫描: {processed}/{total} 文件 ({percentage}%)")

    def display_results(self, file_dict):
        """显示扫描结果"""
        self.tree_files.clear()
        
        for hash_val, files in file_dict.items():
            parent = QTreeWidgetItem(self.tree_files, [f"重复组 ({len(files)}个文件)", "", self.format_size(os.path.getsize(files[0]))])
            for f in files:
                item = QTreeWidgetItem(parent, [os.path.basename(f), f, self.format_size(os.path.getsize(f))])
                item.setCheckState(0, Qt.Unchecked)
        
        # 展开所有项
        self.tree_files.expandAll()
        
        # 启用按钮
        self.btn_delete.setEnabled(True)
        self.btn_select_all.setEnabled(True)
        self.btn_select_duplicates.setEnabled(True)

    def format_size(self, size_bytes):
        """格式化文件大小"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size_bytes < 1024.0:
                return f"{size_bytes:.2f} {unit}"
            size_bytes /= 1024.0
        return f"{size_bytes:.2f} PB"

    def scan_finished(self):
        """扫描完成或被中断"""
        self.btn_select.setText("选择目录")
        self.progress_bar.setValue(100)
        self.status_label.setText("扫描完成")

    def show_preview(self, item):
        if item.childCount() == 0:  # 仅处理子节点
            file_path = item.text(1)
            self.preview_label.clear()
            
            try:
                # 常规图片格式
                if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')):
                    img = Image.open(file_path)
                    img.thumbnail((400, 400))
                    qim = ImageQt.ImageQt(img)
                    self.preview_label.setPixmap(QPixmap.fromImage(qim))
                
                # RAW格式图片
                elif file_path.lower().endswith(('.nef', '.cr2', '.arw', '.cr3', '.orf', '.rw2', '.pef', '.srw', '.dng')):
                    try:
                        with rawpy.imread(file_path) as raw:
                            # 尝试提取缩略图
                            try:
                                thumb = raw.extract_thumb()
                                if thumb.format == rawpy.ThumbFormat.JPEG:
                                    img = Image.open(io.BytesIO(thumb.data))
                                    img.thumbnail((400, 400))
                                    qim = ImageQt.ImageQt(img)
                                    self.preview_label.setPixmap(QPixmap.fromImage(qim))
                                    return
                            except (rawpy.LibRawError, KeyError, ValueError):
                                pass
                                
                            # 如果没有缩略图，则处理完整RAW图像
                            try:
                                rgb = raw.postprocess(use_camera_wb=True, half_size=True, no_auto_bright=True)
                                img = Image.fromarray(rgb)
                                img.thumbnail((400, 400))
                                qim = ImageQt.ImageQt(img)
                                self.preview_label.setPixmap(QPixmap.fromImage(qim))
                            except Exception:
                                self.preview_label.setText("无法预览此RAW文件")
                    except Exception as e:
                        self.preview_label.setText(f"无法预览: {str(e)}")
                else:
                    self.preview_label.setText(f"不支持预览此文件类型\n{os.path.basename(file_path)}")
            except Exception as e:
                self.preview_label.setText(f"预览错误: {str(e)}")

    def select_all_files(self):
        """选择所有文件"""
        for i in range(self.tree_files.topLevelItemCount()):
            parent = self.tree_files.topLevelItem(i)
            for j in range(parent.childCount()):
                item = parent.child(j)
                item.setCheckState(0, Qt.Checked)

    def select_duplicates(self):
        """智能选择重复文件，保留每组的第一个文件"""
        for i in range(self.tree_files.topLevelItemCount()):
            parent = self.tree_files.topLevelItem(i)
            for j in range(parent.childCount()):
                item = parent.child(j)
                # 第一个文件不选，其余全选
                item.setCheckState(0, Qt.Checked if j > 0 else Qt.Unchecked)

    def delete_selected(self):
        """删除选中的文件"""
        # 收集要删除的文件
        files_to_delete = []
        items_to_delete = []
        
        for i in range(self.tree_files.topLevelItemCount()):
            parent = self.tree_files.topLevelItem(i)
            for j in range(parent.childCount()):
                item = parent.child(j)
                if item.checkState(0) == Qt.Checked:
                    files_to_delete.append(item.text(1))
                    items_to_delete.append((parent, item))
        
        if not files_to_delete:
            QMessageBox.information(self, "提示", "请先选择要删除的文件")
            return
        
        # 确认删除
        reply = QMessageBox.question(self, "确认删除", 
                                    f"确定要删除选中的 {len(files_to_delete)} 个文件吗？\n这些文件将被移至回收站。",
                                    QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        
        if reply == QMessageBox.Yes:
            # 执行删除
            error_files = []
            for file_path in files_to_delete:
                try:
                    send2trash.send2trash(file_path)
                except Exception as e:
                    error_files.append((file_path, str(e)))
            
            # 从树中移除已删除的项
            for parent, item in items_to_delete:
                if item.text(1) not in [ef[0] for ef in error_files]:
                    parent.removeChild(item)
            
            # 移除空的父节点
            for i in range(self.tree_files.topLevelItemCount()-1, -1, -1):
                parent = self.tree_files.topLevelItem(i)
                if parent.childCount() == 0:
                    self.tree_files.takeTopLevelItem(i)
            
            # 显示结果
            if error_files:
                error_msg = "以下文件删除失败:\n\n"
                for file_path, error in error_files[:5]:  # 只显示前5个错误
                    error_msg += f"{os.path.basename(file_path)}: {error}\n"
                if len(error_files) > 5:
                    error_msg += f"...以及其他 {len(error_files)-5} 个文件"
                QMessageBox.warning(self, "删除错误", error_msg)
            else:
                QMessageBox.information(self, "删除成功", f"已成功删除 {len(files_to_delete)} 个文件")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = FileDeduplicator()
    ex.show()
    sys.exit(app.exec_())